package org.model.common.interceptor;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.exceptions.ExceptionUtil;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.IsNullExpression;
import net.sf.jsqlparser.expression.operators.relational.NotEqualsTo;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.jdbc.BadSqlGrammarException;

import java.sql.SQLException;
import java.util.List;

/**
 * 全表扫描拦截器.
 *
 * @author Luo
 * @date 2023-06-09 10:09
 */
@Slf4j
public class FullScanInterceptor extends JsqlParserSupport implements InnerInterceptor {

    /**
     * 禁止查询的表
     */
    private final List<String> tables;

    public FullScanInterceptor(List<String> tables) {
        this.tables = tables;
    }

    @Override
    public void beforeQuery(final Executor executor, final MappedStatement ms, final Object parameter, final RowBounds rowBounds, final ResultHandler resultHandler, final BoundSql boundSql) throws SQLException {
        // 获取执行的方法名
        String methodName = ms.getId();
        // 如果是BaseMapper的selectList()方法
        if (methodName.contains(".selectList")) {
            // 校验是否包含limit
            boolean hasLimit = StringUtils.containsIgnoreCase(boundSql.getSql(), "limit");
            if (hasLimit) return;
            // 校验是否包含count(0)
            boolean hasCount = StringUtils.containsIgnoreCase(boundSql.getSql(), "count(0)");
            if (hasCount) return;
            if (InterceptorIgnoreHelper.willIgnoreBlockAttack(ms.getId())) return;
            parserMulti(boundSql.getSql(), null);
        }
    }

    @Override
    protected void processSelect(final Select select, final int index, final String sql, final Object obj) {
        if (null == select.getSelectBody()) {
            return;
        }
        SelectBody selectBody = select.getSelectBody();
        if (selectBody instanceof PlainSelect) {
            PlainSelect plainSelect = (PlainSelect) selectBody;
            Expression where = plainSelect.getWhere();
            FromItem fromItem = plainSelect.getFromItem();
            if (fromItem instanceof Table) {
                Table table = (Table) fromItem;
                // 判断是否带条件查询
                if (fullMatch(where, getTableLogicField(table.getName()))) {
                    if (CollUtil.contains(tables, table.getName())) {
                        log.error("查询表数据时未带任何条件，禁止查询[{}] \n {}", table.getName(), ExceptionUtil.stacktraceToString(new Exception(),5000));
                        throw new BadSqlGrammarException("无参数全表扫描", sql, new SQLException("查询列表数据失败"));
                    }
                    log.warn("查询时未带任何条件：{} \n {}", sql, ExceptionUtil.stacktraceToString(new Exception(),5000));
                }
            }
        }
    }

    /**
     * 匹配是否带条件.
     *
     * @param where where
     * @param logicField logicField
     * @return 结果
     */
    private boolean fullMatch(Expression where, String logicField) {
        if (where == null) {
            return true;
        }
        if (StringUtils.isNotBlank(logicField)) {

            if (where instanceof BinaryExpression) {
                BinaryExpression binaryExpression = (BinaryExpression) where;
                if (StringUtils.equals(binaryExpression.getLeftExpression().toString(), logicField) || StringUtils.equals(binaryExpression.getRightExpression().toString(), logicField)) {
                    return true;
                }
            }

            if (where instanceof IsNullExpression) {
                IsNullExpression binaryExpression = (IsNullExpression) where;
                if (StringUtils.equals(binaryExpression.getLeftExpression().toString(), logicField)) {
                    return true;
                }
            }
        }

        if (where instanceof EqualsTo) {
            // example: 1=1
            EqualsTo equalsTo = (EqualsTo) where;
            return StringUtils.equals(equalsTo.getLeftExpression().toString(), equalsTo.getRightExpression().toString());
        } else if (where instanceof NotEqualsTo) {
            // example: 1 != 2
            NotEqualsTo notEqualsTo = (NotEqualsTo) where;
            return !StringUtils.equals(notEqualsTo.getLeftExpression().toString(), notEqualsTo.getRightExpression().toString());
        } else if (where instanceof OrExpression) {

            OrExpression orExpression = (OrExpression) where;
            return fullMatch(orExpression.getLeftExpression(), logicField) || fullMatch(orExpression.getRightExpression(), logicField);
        } else if (where instanceof AndExpression) {

            AndExpression andExpression = (AndExpression) where;
            return fullMatch(andExpression.getLeftExpression(), logicField) && fullMatch(andExpression.getRightExpression(), logicField);
        } else if (where instanceof Parenthesis) {
            // example: (1 = 1)
            Parenthesis parenthesis = (Parenthesis) where;
            return fullMatch(parenthesis.getExpression(), logicField);
        }

        return false;
    }

    /**
     * 获取表名中的逻辑删除字段
     *
     * @param tableName 表名
     * @return 逻辑删除字段
     */
    private String getTableLogicField(String tableName) {
        if (StringUtils.isBlank(tableName)) {
            return StringPool.EMPTY;
        }

        TableInfo tableInfo = TableInfoHelper.getTableInfo(tableName);
        if (tableInfo == null || !tableInfo.isWithLogicDelete() || tableInfo.getLogicDeleteFieldInfo() == null) {
            return StringPool.EMPTY;
        }
        return tableInfo.getLogicDeleteFieldInfo().getColumn();
    }

}
